#pragma once

#include <SDL.h>
#include <SDL_net.h>
#include <cryptlib.h>
#include <chacha.h>

#include <string>
#include <vector>

class EncryptedStream {
public:
	TCPsocket sock;
	SDLNet_SocketSet sockset;

private:
	CryptoPP::XChaCha20::Encryption enc;
	CryptoPP::XChaCha20::Decryption dec;
	std::vector<char> buffer;

public:
	EncryptedStream();
	~EncryptedStream();

	void create(TCPsocket sock_);
	void destroy();

	int send(const void *data, int len);
	int receive(void *data, int len, bool exact);
	void dropIncoming();

	bool check(int timeout);

	template <class S>
	bool sendString(const S& s, int maxlensize) {
		size_t len = s.size();
		size_t maxlen = (((size_t)1) << (8 * maxlensize)) - 1;
		if (len >= maxlen) len = maxlen;

		if (send(&len, maxlensize) != maxlensize) return false;

		if (len > 0) {
			size_t len2 = len * sizeof(S::value_type);
			if (send(s.c_str(), len2) != len2) return false;
		}

		return true;
	}

	template <class S>
	bool receiveString(S& s, int maxlensize) {
		s.clear();

		size_t len = 0;
		if (receive(&len, maxlensize, true) != maxlensize) return false;

		if (len > 0) {
			std::vector<S::value_type> buf(len);
			size_t len2 = len * sizeof(S::value_type);
			if (receive(buf.data(), len2, true) != len2) return false;
			s = S(buf.data(), len);
		}

		return true;
	}
};
